library(tidyverse)
library(ggpubr)
library(LDlinkR)
library(ggrepel)

# plot 1kgp population on globe -------------
# Import the data with coordinates (download from https://www.internationalgenome.org/data-portal/population)
world <- map_data("world")

geography_1kgp <- read_delim('Archive/covid19/ref/igsr_populations.tsv',
                             name_repair = make.names) |>
  filter(!is.na(Population.code) & !is.na(Superpopulation.code))

geography_1kgp |>
  glimpse()

# Plot the map --------------- 
# group = group connects the points in the correct order
ggplot(data = world, aes(x = long, y = lat, group = group)) + 
  geom_polygon(fill = 'white', color = 'black') +
  geom_point(data = geography_1kgp,
             aes(x = Population.longitude,
                 y = Population.latitude,
                 color = Superpopulation.code),
             size = 5,
             inherit.aes = FALSE)

# try plot from 20W to 20W
# a function to correct polygons
correct_polygon <- function(tbl, group_id){
  east_part <- tbl |>
    filter(long >= -20) |>
    arrange(order) |>
    mutate(order = seq_along(long)) |>
    add_column(group = group_id)
  
  tbl |>
    filter(long < -20) |>
    arrange(order) |>
    mutate(order = seq_along(long)) |>
    add_column(group = group_id + 2000) |>
    bind_rows(east_part)
}

bad_group <- world |>
  group_by(group) |>
  summarise(cross0 = min(long + 20) * max(long + 20)) |>
  filter(cross0 < 0) |>
  pull(group)

corrected_group <- world |>
  filter(group %in% bad_group) |>
  nest(.by = group) |>
  rowwise() |>
  mutate(corrected = list(correct_polygon(data, group))) |>
  pull(corrected) |>
  reduce(bind_rows)

new_world <- world |>
  filter(!(group %in% bad_group)) |>
  bind_rows(corrected_group) |>
  mutate(long = case_when(long < -20 ~ long + 360,
                          .default = long))

ggplot(data = new_world, aes(x = long, y = lat, group = group)) + 
  geom_polygon(fill = 'white', color = 'black') +
  coord_sf(ylim = c(-55, 80))

# show IGHG1-GR freq on world map ----------
ighg1_block_1kgp <- read_delim('Archive/covid19/results/1KGP_subpop_MAF.csv')

g2r_position <- 106204113

g2r_geograhy <- ighg1_block_1kgp |>
  filter(POS == g2r_position) |>
  left_join(geography_1kgp, join_by(pop == Population.code))

ggplot(data = world, aes(x = long, y = lat, group = group)) + 
  geom_polygon(fill = 'white', color = 'black') +
  geom_point(data = g2r_geograhy,
             aes(x = Population.longitude,
                 y = Population.latitude,
                 color = MAF),
             size = 5,
             inherit.aes = FALSE) +
  scale_color_gradient2(midpoint = .5)+
  coord_sf(ylim = c(-55, 80))

ggplot(data = new_world, aes(x = long, y = lat, group = group)) + 
  geom_polygon(fill = 'white', color = 'black') +
  geom_point(data = g2r_geograhy,
             aes(x = case_when(Population.longitude < -20 ~ Population.longitude + 360,
                               .default = Population.longitude),
                 y = Population.latitude,
                 color = MAF),
             size = 5,
             inherit.aes = FALSE) +
  scale_color_gradient2(midpoint = .5) +
  coord_sf(ylim = c(-55, 80))

# simons dgp --------
ighg1_block_sdgp <- read_delim('Archive/covid19/data/sdgp14-igg1-block.vcf.gz', comment = '##') |>
  select(-c(ID, QUAL, FILTER, FORMAT)) |>
  mutate(INFO = str_remove_all(INFO, '[A-Z]+=')) |>
  separate_wider_delim(INFO, names = c('MAF_ref_pop',
                          'MAF_target',
                          'imputation_info',
                          'allele_count_in_genotype',
                          'total_allele_num_in_genotype'), delim = ';')

ighg1_block_sdgp

g2r_simons <- ighg1_block_sdgp |>
  filter(POS == g2r_position) |>
  pivot_longer(cols = 10:last_col()) |>
  separate_wider_delim(value, names = c('unphased_genotype',
                                        'genotype_dosage',
                                        'genotype_posterior',
                                        'haplotype_sampled',
                                        'Phred_likelihood',
                                        'allelic_depth'),
                       delim = ':') |>
  mutate(Population.name = str_extract(name, '(?<=_).+(?=-)')) |>
  group_by(Population.name) |>
  summarize(maf = mean(as.numeric(genotype_dosage)))

meta_simons <- read_delim('Archive/covid19/ref/igsr_populations.tsv',
                             name_repair = make.names) |>
  filter(str_detect(Data.collections, 'Simons'))

geo_simons <- g2r_simons |>
  left_join(meta_simons, multiple = 'all')

ggplot(data = world, aes(x = long, y = lat, group = group)) + 
  geom_polygon(fill = 'white', color = 'black') +
  geom_point(data = geo_simons,
             aes(x = Population.longitude,
                 y = Population.latitude,
                 color = maf),
             size = 5,
             inherit.aes = FALSE)

# function to get sub-maf of snp from 1kgp -------
subpop_1kgp <- list_pop() |>
  filter(str_detect(pop_name, '[a-z]'))

get_one_subpop <- function(x, snp_list){
  SNPclip(snp_list,
          token = Sys.getenv('LDLINK_TOKEN'),
          pop = x) |>
    add_column(pop = x)
}

safely_get_one_pop <- safely(get_one_subpop)

safely_get_one_pop(x = 'CDX', snp_list = c('rs333','rs1800940'))

list_sub_maf_1kgp <- function(snps){
  subpop_1kgp$pop_code |>
    map(\(x)safely_get_one_pop(x = x, snp_list = snps),
        .progress = TRUE) |>
    list_flatten() |>
    keep(is.data.frame) |>
    list_rbind() |>
    as_tibble()
}

li_guo_aid <- read_delim('Archive/covid19/data/Li-Guo_AID.txt', name_repair = make.names)

li_guo_avail <- li_guo_aid |>
  filter(str_detect(variant, 'rs')) |>
  separate_wider_delim(loci..GRCh37., names = c('chr','pos'), delim = ':') |>
  mutate(pos = as.numeric(pos),
         chr = str_glue('chr{chr}'),
         variant = str_remove(variant, '\\s.+'))

safe_list_1kgp <- safely(list_sub_maf_1kgp)

li_guo_1kgp <- li_guo_avail |>
  group_by(chr) |>
  summarise(rs_list = list(variant)) |>
  pull(rs_list) |>
  map(\(x)append(x, 'rs1')) |>
  map(safe_list_1kgp)

li_guo_1kgp |>
  list_flatten() |>
  keep(is.data.frame) |>
  list_rbind() |>
  as_tibble() |>
  filter(RS_Number != 'rs1') |>
  separate(Alleles, into = c('ref','alt'), sep = ',') |>
  mutate(maf = str_extract(alt, '(?<==).+') |> as.numeric()) |>
  select(c(RS_Number, pop, maf)) |>
  left_join(subpop_1kgp, join_by(pop == pop_code)) |>
  filter(!is.na(super_pop_code)) ->
  li_guo_res

li_guo_res |>
  mutate(superpop.index = fct(super_pop_code) |> as.numeric()) |>
  ggplot(aes(fct_reorder(pop, superpop.index), maf, fill = super_pop_code)) +
  geom_col() +
  facet_wrap(~RS_Number) +
  theme_pubr() +
  labs_pubr() +
  theme(axis.text.x = element_blank()) +
  labs(x = 'Population', y = 'MAF')
    
# CCR5-del32 rs333 --------
rs333 <- list_sub_maf_1kgp(c('rs333','rs1800940'))

rs333_geo <- rs333 |>
  as_tibble() |>
  filter(RS_Number == 'rs333') |>
  mutate(maf = str_extract(Alleles, '(?<=-=).+') |> as.double(),
         Position = NULL, Details = NULL, Alleles = NULL,
         Population.code = pop) |>
  left_join(geography_1kgp) |>
  select(Population.code, maf ,Superpopulation.code,
         Population.longitude, Population.latitude) |>
  filter(!(Population.code %in% c('ITU','STU')))

world |>
  ggplot(aes(x = long, y = lat, group = group)) + 
  geom_polygon(fill = 'white', color = 'grey') +
  geom_point(data = rs333_geo,
             aes(x = Population.longitude,
                 y = Population.latitude,
                 color = maf),
             size = 3,
             inherit.aes = FALSE) +
  ggrepel::geom_text_repel(data = filter(rs333_geo, maf > 0),
                           aes(x = Population.longitude,
                               y = Population.latitude,
                               label = str_c(maf * 100, '%')),
                           nudge_y = 5,
                           inherit.aes = FALSE) +
  theme_bw() +
  scale_color_viridis_c() +
  coord_cartesian(xlim = c(-160,170), ylim = c(-50,80)) +
  labs(x = 'longitude', y = 'latitude',
       color = 'Minor allele freq', title = 'CCR5-del32 (rs333) in 1KGP data')

ccr5_unip <- read_delim('https://www.ebi.ac.uk/proteins/api/variation/P51681?format=gff',
           comment = '#', col_names = F)

ccr5.clinvar <- ccr5_unip |>
  filter(str_detect(X9, 'ClinVar'), str_detect(X9, 'rs')) |>
  select(X9) |>
  mutate(id = str_extract(X9, 'rs\\d+'))

ccr5.ldtrait <- ccr5.clinvar$id |>
  LDtrait(token = Sys.getenv('LDLINK_TOKEN'))

ccr5.disgen <- read_tsv('Archive/covid19/ref/vda__source-all.tsv') |>
  filter(Variant != 'rs333')

ccr5.maf <- ccr5.disgen |>
  select(Variant, Disease) |>
  filter(str_detect(Disease, 'VIRUS')) |>
  pull(Variant) |>
  unique() |>
  c('rs1799988') |>
  list_sub_maf_1kgp()

ccr5.maf <- ccr5.maf |>
  mutate(maf = str_extract(Alleles, ',.+') |> str_extract('.\\.\\d+') |> as.double(),
         Population.code = pop) |>
  select(RS_Number, Population.code, maf) |>
  left_join(geography_1kgp) |>
  select(Population.code, maf ,Superpopulation.code, RS_Number,
         Population.longitude, Population.latitude) |>
  filter(!(Population.code %in% c('ITU','STU')))

ccr5.maf |>
  ggplot(aes(Population.code, maf)) +
  geom_col() +
  facet_wrap(~RS_Number, scales = 'free_y')

plot_1kgp_globe <- function(df, snp){
  world |>
  ggplot(aes(x = long, y = lat, group = group)) + 
  geom_polygon(fill = 'white', color = 'grey') +
  geom_point(data = df,
             aes(x = Population.longitude,
                 y = Population.latitude,
                 color = maf),
             size = 3,
             inherit.aes = FALSE) +
  geom_text_repel(data = filter(df, maf > 0),
                  aes(x = Population.longitude,
                      y = Population.latitude,
                      label = str_c(maf * 100, '%')),
                  nudge_y = 5,
                  inherit.aes = FALSE) +
  theme_bw() +
  scale_color_viridis_c(limits = c(0,NA)) +
  coord_cartesian(xlim = c(-160,170), ylim = c(-50,80)) +
  labs(x = 'longitude', y = 'latitude',
       color = 'Minor allele freq', title = str_glue('{snp} in 1KGP data'))}

ccr5.maf |>
  filter(str_detect(RS_Number, '9420')) |>
  plot_1kgp_globe('CCR5-S63C')

ccr5.maf |>
  filter(str_detect(RS_Number, '0560')) |>
  plot_1kgp_globe('CCR5-C101* (rs1800560)')

ccr5.maf |>
  filter(str_detect(RS_Number, 'rs1799987')) |>
  plot_1kgp_globe('rs1799987')

ccr5.maf |>
  filter(str_detect(RS_Number, 'rs41469351')) |>
  plot_1kgp_globe('rs41469351')

old.snp <- list_sub_maf_1kgp('rs333') 

old.snp |> 
  mutate(maf = str_extract(Alleles, ',.+') |> str_extract('.\\.\\d+') |> as.double(),
         Population.code = pop) |>
  select(RS_Number, Population.code, maf) |>
  left_join(geography_1kgp) |>
  select(Population.code, maf ,Superpopulation.code, RS_Number,
         Population.longitude, Population.latitude) |>
  filter(!(Population.code %in% c('ITU','STU')))

ccr5.maf |>
  filter(str_detect(RS_Number, 'rs1799988')) |>
  plot_1kgp_globe('rs1799988')

# GSDMB rs2123685 ----
gsdmb <- list_sub_maf_1kgp(c('rs2123685','rs12453507'))

gsdmb |>
  filter(str_ends(RS_Number, '85')) |>
  mutate(maf = str_extract(Alleles, ',.+') |> str_extract('.\\.\\d+') |> as.double(),
         Population.code = pop) |>
  select(RS_Number, Population.code, maf) |>
  left_join(geography_1kgp) |>
  select(Population.code, maf ,Superpopulation.code, RS_Number,
         Population.longitude, Population.latitude) |>
  filter(!(Population.code %in% c('ITU','STU'))) |>
  plot_1kgp_globe('rs2123685')

  
